import io
from typing import Annotated

from PIL import Image
from typing_extensions import TypedDict

from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages

import llms

"""
定义图形时要做的第一件事就是定义图形。它由图形的架构以及指定如何对状态应用更新的 reducer 函数组成。
在我们的示例中是一个具有单个键的键：。该键使用 add_messages reducer 函数进行注释，该函数告诉 LangGraph 将新消息追加到现有列表，而不是覆盖它。没有注释的状态键将被每次更新覆盖，并存储最新的值。查看此概念指南，了解有关状态、化简器和其他低级概念的更多信息。
"""
class State(TypedDict):
    # Messages have the type "list". The `add_messages` function
    # in the annotation defines how this state key should be updated
    # (in this case, it appends messages to the list, rather than overwriting them)
    messages: Annotated[list, add_messages]

# 构建一个图
graph_builder = StateGraph(State)


llm = llms.getLLM()


def chatbot(state: State):
    return {"messages": [llm.invoke(state["messages"])]}


# The first argument is the unique node name
# The second argument is the function or object that will be called whenever
# the node is used.
# 将节点添加到图
graph_builder.add_node("chatbot", chatbot)

# 从哪里开始工作
graph_builder.add_edge(START, "chatbot")
# 从哪里结束
graph_builder.add_edge("chatbot", END)
# 编译图
graph = graph_builder.compile()

# 绘制图形成图片
# im=Image.open(io.BytesIO(graph.get_graph().draw_mermaid_png()))
# # im.save("graph.png")
# im.show()

while True:
    user_input = input("User: ")
    if user_input.lower() in ["quit", "exit", "q"]:
        print("Goodbye!")
        break
    for event in graph.stream({"messages": ("user", user_input)}):
        for value in event.values():
            print("Assistant:", value["messages"][-1].content)